/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.web.service;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.core.execution.JobClient;
import org.apache.flink.table.api.TableResult;
import org.apache.flink.table.catalog.Column;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.client.cli.CliStrings;
import org.apache.flink.table.client.cli.CliUtils;
import org.apache.flink.table.client.config.SqlClientOptions;
import org.apache.flink.table.client.gateway.Executor;
import org.apache.flink.table.client.gateway.SqlExecutionException;
import org.apache.flink.table.operations.*;
import org.apache.flink.table.operations.command.*;
import org.apache.flink.table.utils.PrintUtils;
import org.apache.flink.types.Row;
import org.apache.flink.web.domain.SqlResult;
import org.jline.utils.AttributedString;
import org.jline.utils.AttributedStringBuilder;
import org.jline.utils.AttributedStyle;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.ZoneId;
import java.util.*;

import static org.apache.flink.table.api.config.TableConfigOptions.TABLE_DML_SYNC;
import static org.apache.flink.table.api.internal.TableResultImpl.TABLE_RESULT_OK;
import static org.apache.flink.table.client.cli.CliStrings.*;
import static org.apache.flink.table.client.config.YamlConfigUtils.*;
import static org.apache.flink.util.Preconditions.checkState;

/** SQL CLI client. */
public class SqlSession implements AutoCloseable {

    private static final Logger LOG = LoggerFactory.getLogger(SqlSession.class);

    private final Executor executor;

    private final String sessionId;

    private final String prompt;

    private Integer maxRows = Integer.MAX_VALUE;

    private boolean isStatementSetMode;

    private List<ModifyOperation> statementSetOperations;

    /**
     * Creates a CLI instance with a custom terminal. Make sure to close the CLI instance afterwards
     * using {@link #close()}.
     */
    @VisibleForTesting
    public SqlSession(
            String sessionId,
            Executor executor,
            Integer maxRows) {
        this.sessionId = sessionId;
        this.executor = executor;
        if (maxRows != null) {
            this.maxRows = maxRows;
        }

        // create prompt
        prompt = new AttributedStringBuilder()
                        .style(AttributedStyle.DEFAULT.foreground(AttributedStyle.GREEN))
                        .append("Flink SQL")
                        .style(AttributedStyle.DEFAULT)
                        .append("> ")
                        .toAnsi();
    }

    public String getSessionId() {
        return this.sessionId;
    }

    public Executor getExecutor() {
        return executor;
    }

    /** Closes the CLI instance. */
    public void close() {
    }

    // --------------------------------------------------------------------------------------------


    /**
     * Execute content from Sql file and prints status information and/or errors on the terminal.
     *
     * @param content SQL file content
     */
    public SqlResult execute(String content) {
        LOG.info(CliStrings.messageInfo(CliStrings.MESSAGE_EXECUTE_FILE).toAnsi());

        SqlResult sqlResult = new SqlResult(sessionId);
        for (String statement : splitSql(content)) {
            LOG.info(new AttributedString(String.format("%s%s", prompt, statement)).toString());
            // cancel execution when meet error or ctrl + C;
            try {
                sqlResult.addJob(callOperation(executor.parseStatement(sessionId, statement)));
            } catch (SqlExecutionException e) {
                printExecutionException(e);
                sqlResult.failed(e);
                return sqlResult;
            }
        }
        return sqlResult;
    }

    private SqlResult.JobResult callOperation(Operation operation) {
        if (operation instanceof QuitOperation) {
            // QUIT/EXIT
            return null;
        } else if (operation instanceof ClearOperation) {
            // CLEAR
            return null;
        } else if (operation instanceof HelpOperation) {
            // HELP
            return null;
        } else if (operation instanceof SetOperation) {
            // SET
            return callSet((SetOperation) operation);
        } else if (operation instanceof ResetOperation) {
            // RESET
            return callReset((ResetOperation) operation);
        } else if (operation instanceof CatalogSinkModifyOperation) {
            // INSERT INTO/OVERWRITE
            return callInsert((CatalogSinkModifyOperation) operation);
        } else if (operation instanceof QueryOperation) {
            // SELECT
            return callSelect((QueryOperation) operation);
        } else if (operation instanceof ExplainOperation) {
            // EXPLAIN
            return callExplain((ExplainOperation) operation);
        } else if (operation instanceof BeginStatementSetOperation) {
            // BEGIN STATEMENT SET
            return callBeginStatementSet();
        } else if (operation instanceof EndStatementSetOperation) {
            // END
            return callEndStatementSet();
        } else {
            // fallback to default implementation
            return executeOperation(operation);
        }
    }


    private SqlResult.JobResult callReset(ResetOperation resetOperation) {
        // reset all session properties
        if (!resetOperation.getKey().isPresent()) {
            executor.resetSessionProperties(sessionId);
            printInfo(CliStrings.MESSAGE_RESET);
        }
        // reset a session property
        else {
            String key = resetOperation.getKey().get();
            executor.resetSessionProperty(sessionId, key);
            printSetResetConfigKeyMessage(key, MESSAGE_RESET_KEY);
        }
        return null;
    }

    private SqlResult.JobResult callSet(SetOperation setOperation) {
        // set a property
        if (setOperation.getKey().isPresent() && setOperation.getValue().isPresent()) {
            String key = setOperation.getKey().get().trim();
            String value = setOperation.getValue().get().trim();
            executor.setSessionProperty(sessionId, key, value);
            printSetResetConfigKeyMessage(key, MESSAGE_SET_KEY);
        }
        // show all properties
        else {
            final Map<String, String> properties = executor.getSessionConfigMap(sessionId);
            if (properties.isEmpty()) {
                LOG.info(CliStrings.messageInfo(CliStrings.MESSAGE_EMPTY).toAnsi());
            } else {
                List<String> prettyEntries = getPropertiesInPretty(properties);
                prettyEntries.forEach(LOG::info);
            }
        }
        return null;
    }

    private SqlResult.JobResult callSelect(QueryOperation operation) {
        return executeOperation(operation);
    }


    private SqlResult.JobResult callInsert(CatalogSinkModifyOperation operation) {
        if (isStatementSetMode) {
            statementSetOperations.add(operation);
            printInfo(CliStrings.MESSAGE_ADD_STATEMENT_TO_STATEMENT_SET);
            return null;
        } else {
            return callInserts(Collections.singletonList(operation));
        }
    }

    private SqlResult.JobResult callInserts(List<ModifyOperation> operations) {
        printInfo(CliStrings.MESSAGE_SUBMITTING_STATEMENT);

        boolean sync = executor.getSessionConfig(sessionId).get(TABLE_DML_SYNC);
        if (sync) {
            printInfo(MESSAGE_WAIT_EXECUTE);
        }
        TableResult tableResult = executor.executeModifyOperations(sessionId, operations);
        checkState(tableResult.getJobClient().isPresent());

        if (sync) {
            LOG.info(CliStrings.messageInfo(MESSAGE_FINISH_STATEMENT).toAnsi());
        } else {
            LOG.info(CliStrings.messageInfo(MESSAGE_STATEMENT_SUBMITTED).toAnsi());
            LOG.info(String.format("Job ID: %s\n", tableResult.getJobClient().get().getJobID().toString()));
        }
        return SqlResult.newJobResult(tableResult.getJobClient().get().getJobID().toString());
    }

    public SqlResult.JobResult callExplain(ExplainOperation operation) {
        TableResult tableResult = executor.executeOperation(sessionId, operation);
        // show raw content instead of tableau style
        final String explanation =
                Objects.requireNonNull(tableResult.collect().next().getField(0)).toString();
        LOG.info(explanation);
        return null;
    }

    private SqlResult.JobResult callBeginStatementSet() {
        isStatementSetMode = true;
        statementSetOperations = new ArrayList<>();
        printInfo(CliStrings.MESSAGE_BEGIN_STATEMENT_SET);
        return null;
    }

    private SqlResult.JobResult callEndStatementSet() {
        if (isStatementSetMode) {
            isStatementSetMode = false;
            SqlResult.JobResult jobResult = null;
            if (!statementSetOperations.isEmpty()) {
                jobResult = callInserts(statementSetOperations);
            } else {
                printInfo(CliStrings.MESSAGE_NO_STATEMENT_IN_STATEMENT_SET);
            }
            statementSetOperations = null;
            return jobResult;
        } else {
            throw new SqlExecutionException(MESSAGE_STATEMENT_SET_END_CALL_ERROR);
        }
    }

    private SqlResult.JobResult executeOperation(Operation operation) {
        TableResult result = executor.executeOperation(sessionId, operation);
        if (TABLE_RESULT_OK == result) {
            // print more meaningful message than tableau OK result
            printInfo(MESSAGE_EXECUTE_STATEMENT);
            return null;
        } else {
            Optional<JobClient> jobClientOptional = result.getJobClient();
            if(jobClientOptional.isPresent()) {
                JobClient jobClient = jobClientOptional.get();
                ZoneId timeZone = CliUtils.getSessionTimeZone(executor.getSessionConfig(sessionId));
                return getJobResult(jobClient.getJobID().toString(), result.getResolvedSchema(), result.collect(), timeZone);
            } else {
                return null;
            }
        }
    }

    // --------------------------------------------------------------------------------------------

    private void printExecutionException(Throwable t) {
        final String errorMessage = CliStrings.MESSAGE_SQL_EXECUTION_ERROR;
        LOG.warn(errorMessage, t);
        boolean isVerbose = executor.getSessionConfig(sessionId).get(SqlClientOptions.VERBOSE);
        LOG.info(CliStrings.messageError(errorMessage, t, isVerbose).toAnsi());
    }

    private void printInfo(String message) {
        LOG.info(CliStrings.messageInfo(message).toAnsi());
    }

    private void printWarning(String message) {
        LOG.info(CliStrings.messageWarning(message).toAnsi());
    }

    private void printSetResetConfigKeyMessage(String key, String message) {
        boolean isRemovedKey = isRemovedKey(key);
        boolean isDeprecatedKey = isDeprecatedKey(key);

        // print warning information if the given key is removed or deprecated
        if (isRemovedKey || isDeprecatedKey) {
            String warningMsg =
                    isRemovedKey
                            ? MESSAGE_REMOVED_KEY
                            : String.format(
                                    MESSAGE_DEPRECATED_KEY,
                                    key,
                                    getOptionNameWithDeprecatedKey(key));
            printWarning(warningMsg);
        }

        // when the key is not removed, need to print normal message
        if (!isRemovedKey) {
            LOG.info(CliStrings.messageInfo(message).toAnsi());
        }
    }


    /**
     * 分割SQL语句, 删除注释
     *
     * @param allSql sql
     * @return sqls
     */
    public static List<String> splitSql(String allSql) {
        return splitSql(allSql, true);
    }

    /**
     * 分割SQL语句
     *
     * @param allSql  sql
     * @param rmNotes 删除注释
     * @return sqls
     */
    public static List<String> splitSql(String allSql, boolean rmNotes) {
        List<String> results = new ArrayList<>();

        int i = 0;
        while (i < allSql.length()) {
            char c = allSql.charAt(i);
            if (c == ';') {
                // 切分段落
                String item = allSql.substring(0, i).trim();
                if (!item.isEmpty()) {
                    results.add(item);
                }
                allSql = allSql.substring(i + 1);
                i = 0;
                continue;
            } else if (c == '\'') {
                // 字符包裹, 跳过内容解析
                i = getIndex(allSql, i + 1, '\'');
            } else if (c == '"') {
                // 字符串包裹, 跳过内容解析
                i = getIndex(allSql, i + 1, '"');
            } else if (c == '`') {
                // 字段包裹, 跳过内容解析
                i = getIndex(allSql, i + 1, '`');
            } else if (c == '/' && allSql.charAt(i + 1) == '*') {
                // 多行注释
                int wrapEnd = getIndex(allSql, i + 2, "*/");
                if (rmNotes) {
                    //删除注释
                    allSql = allSql.substring(0, i) + allSql.substring(wrapEnd + 2);
                    continue;
                } else {
                    //跳过内容解析
                    i = wrapEnd + 1;
                }
            } else if (c == '-' && allSql.charAt(i + 1) == '-') {
                // 单行注释
                int wrapEnd = allSql.indexOf("\n", i);
                if (rmNotes) {
                    //删除注释
                    allSql = wrapEnd == -1 ? allSql.substring(0, i) : (allSql.substring(0, i) + "\n" + allSql.substring(wrapEnd + 1));
                    continue;
                } else {
                    //跳过内容解析
                    i = wrapEnd == -1 ? allSql.length() : wrapEnd;
                }
            }
            i++;
        }

        String item = allSql.trim();
        if (!item.isEmpty()) {
            results.add(item);
        }
        return results;
    }

    /**
     * 获取字符串位置(找不到时异常)
     *
     * @param eval       字符串
     * @param beginIndex 起始位置
     * @return 字符串位置
     */
    private static int getIndex(String eval, int beginIndex, String str) {
        int index = getIndex(eval, beginIndex, str.charAt(0));
        if (str.length() == 1) {
            return index;
        }
        if (eval.length() < index + str.length()) {
            throw new RuntimeException("语法错误");
        }
        for (int i = 1; i < str.length(); i++) {
            if (eval.charAt(index + i) != str.charAt(i)) {
                return getIndex(eval, index + 1, str);
            }
        }
        return index;
    }

    /**
     * 获取字符串位置(找不到时异常)
     *
     * @param eval       字符串
     * @param beginIndex 起始位置
     * @return 字符串位置
     */
    private static int getIndex(String eval, int beginIndex, char c) {
        for (int i = beginIndex; i < eval.length(); i++) {
            char ch = eval.charAt(i);
            if (ch == c) {
                return i;
            }
        }
        throw new RuntimeException("语法错误");
    }


    /**
     * 获取执行结果
     * @param jobId jobId
     * @param resolvedSchema schema
     * @param it 返回结果
     * @param timeZone 时区
     * @return 执行结果
     */
    private SqlResult.JobResult getJobResult(
            String jobId,
            ResolvedSchema resolvedSchema,
            Iterator<Row> it,
            ZoneId timeZone) {

        // columnNames
        final List<Column> columns = resolvedSchema.getColumns();
        SqlResult.ResultColumn[] resultColumns = columns.stream().map(SqlResult.ResultColumn::new).toArray(SqlResult.ResultColumn[]::new);

        // result
        List<String[]> result = new ArrayList<>();
        while (it.hasNext() && result.size()< maxRows) {
            result.add(PrintUtils.rowToString(it.next(), PrintUtils.NULL_COLUMN, false, resolvedSchema, timeZone));
        }
        return SqlResult.newJobResult(jobId, resultColumns, result);
    }

    public static class EmbeddedShutdownThread extends Thread {

        private final String sessionId;
        private final Executor executor;

        public EmbeddedShutdownThread(String sessionId, Executor executor) {
            this.sessionId = sessionId;
            this.executor = executor;
        }

        @Override
        public void run() {
            // Shutdown the executor
            System.out.println("\nShutting down the session...");
            executor.closeSession(sessionId);
            System.out.println("done.");
        }
    }
}
